# web_1px_pen_pinecone.py  ——  使用Pinecone云服务进行推理的手写数字识别
from __future__ import annotations
import os
import logging
from typing import Any, Tuple
import numpy as np
from collections import Counter

try:
    import gradio as gr
except Exception:
    gr = None

try:
    from PIL import Image
except Exception:
    Image = None

# 初始化日志配置
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 初始化Pinecone客户端
try:
    from pinecone import Pinecone
    pinecone = Pinecone(api_key="pcsk_4zcKEp_P5uHzGJS6Eq9UhAcqWoiFv6seAxnnxPhFb8yDkHykkHKsCNAiBMvr1TH9Nkzz3f")
    index_name = "mnist-index"
    index = pinecone.Index(index_name)
    logging.info(f"成功连接到Pinecone索引 '{index_name}'")
    PINECONE_AVAILABLE = True
except Exception as e:
    logging.error(f"Pinecone初始化失败: {e}")
    PINECONE_AVAILABLE = False


def preprocess_image_for_digits(img: Any) -> Tuple[np.ndarray, np.ndarray]:
    if img is None:
        raise ValueError("输入图像为空")
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype(np.uint8))
    img = img.convert("L").resize((8, 8), Image.LANCZOS)
    arr = np.asarray(img, dtype=np.float32)
    arr = 255.0 - arr
    arr = (arr / 255.0) * 16.0
    vec = arr.flatten()
    processed_img = (vec.reshape(8, 8) / 16.0 * 255).astype(np.uint8)
    Image.fromarray(processed_img).save("debug_processed.png")
    logging.debug(f"64维向量: {np.round(vec, 2)}")
    return vec, processed_img


def predict_from_canvas(img: Any) -> Tuple[str, np.ndarray]:
    try:
        if not PINECONE_AVAILABLE:
            return "错误：无法连接到Pinecone服务", np.zeros((8, 8), dtype=np.uint8)
        
        x, processed_img = preprocess_image_for_digits(img)
        
        # 使用Pinecone进行查询
        query_data = x.tolist()
        results = index.query(
            vector=query_data,
            top_k=11,  # 与训练时保持一致
            include_metadata=True
        )
        
        # 提取标签并进行投票
        labels = [match['metadata']['label'] for match in results['matches']]
        final_prediction = Counter(labels).most_common(1)[0][0]
        
        # 显示前几个最相似的结果
        top_matches_info = "\n\nTop 匹配结果:\n" + "\n".join([
            f"{i+1}. 数字 {match['metadata']['label']}: 距离 {match['score']:.4f}" 
            for i, match in enumerate(results['matches'][:5])
        ])
        
        result = f"预测结果：{final_prediction}{top_matches_info}"
        logging.info(f"预测结果: {final_prediction}")
        return result, processed_img
    except Exception as e:
        logging.error(f"预测错误: {str(e)}")
        return f"预测错误：{str(e)}", np.zeros((8, 8), dtype=np.uint8)


def clear_canvas() -> Tuple[None, str, np.ndarray]:
    """Clear the canvas and reset outputs"""
    return None, "请在画布上绘制一个数字", np.zeros((8, 8), dtype=np.uint8)


def create_and_launch_app():
    if gr is None:
        raise RuntimeError("请 pip install gradio==3.50.2")

    with gr.Blocks(title="KNN 手写数字识别 (Pinecone版)") as iface:
        gr.Markdown("# KNN 手写数字识别（Pinecone云服务版）")
        gr.Markdown("在画布上绘制一个数字，然后点击 Submit 查看预测结果（使用Pinecone云服务进行推理）")
        
        with gr.Row():
            with gr.Column():
                canvas = gr.Image(
                    source="canvas",
                    type="numpy",
                    image_mode="L",
                    shape=(40, 40),
                    label="在白画布内写一个小而细的数字（0-9）"
                )
                with gr.Row():
                    submit_btn = gr.Button("Submit", variant="primary")
                    clear_btn = gr.Button("Clear")
            
            with gr.Column():
                output_text = gr.Textbox(label="预测结果", value="请在画布上绘制一个数字")
                processed_image = gr.Image(label="预处理后的 8×8 图像", shape=(160, 160))
        
        gr.Markdown("### 提示")
        gr.Markdown("- debug_processed.png 文件会保存预处理后的图像\n- 绘制时请尽量在画布中央书写数字\n- 每次推理都使用Pinecone云服务")
        
        # Event handlers
        submit_btn.click(
            fn=predict_from_canvas,
            inputs=canvas,
            outputs=[output_text, processed_image]
        )
        
        clear_btn.click(
            fn=clear_canvas,
            inputs=None,
            outputs=[canvas, output_text, processed_image]
        )
        
        # Also clear on image change (optional)
        canvas.change(
            fn=lambda: ("请在画布上绘制一个数字", np.zeros((8, 8), dtype=np.uint8)),
            inputs=None,
            outputs=[output_text, processed_image]
        )
    
    iface.launch(share=False, server_name="127.0.0.1", server_port=7860)


if __name__ == "__main__":
    create_and_launch_app()